package com.tuniu.agents.impl.trip;

import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.agent.AgentManager;
import com.tuniu.agents.common.handler.IStreamResponseHandler;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.impl.trip.enums.PlannerEnum;
import com.tuniu.agents.memory.MemoryManager;
import com.tuniu.agents.message.MessageSender;
import com.tuniu.agents.util.I18nUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import java.util.Collections;
import java.util.Map;

@Slf4j
public class PlannerStreamResponseHandler implements IStreamResponseHandler {


    private final AgentManager agentManager;
    private final ChatEnvironmentMemory chatEnvironmentMemory;
    private final MessageSender messageSender;
    private final MemoryManager memoryManager;

    private static final Logger logger = LoggerFactory.getLogger(PlannerStreamResponseHandler.class);
    
    // ID of the Summarize Agent
    private static final String SUMMARIZE_AGENT_ID = "SummarizeAgent";
    // ID of the Coordinator Agent
    private static final String COORDINATOR_AGENT_ID = "Coordinator";

    public PlannerStreamResponseHandler(ChatEnvironmentMemory chatEnvironmentMemory, AgentManager agentManager, MessageSender messageSender, MemoryManager memoryManager) {
        this.chatEnvironmentMemory = chatEnvironmentMemory;
        this.agentManager = agentManager;
        this.messageSender = messageSender;
        this.memoryManager = memoryManager;
    }

    public Flux<String> handle(Map<String, Object> handlerContext, Flux<String> flux) {
        String conversationId = (String) handlerContext.get("conversationId");
        String plannerOutMessage = I18nUtil.getMessage(conversationId, PlannerEnum.PLANNER_OUTPUT.getMessageKey());
        messageSender.send(conversationId, BaseAgentNewResp.sendAnalysisMsg(plannerOutMessage, 1, null));

        String message = (String) handlerContext.get("message");
        StringBuilder contentBuffer = new StringBuilder();
        
        Flux<String> originalResponse = flux.flatMap(chunk -> {
            contentBuffer.append(chunk);
            if(org.apache.commons.lang3.StringUtils.isNotBlank(chunk)){
                return Flux.just(BaseAgentNewResp.sendAnswerMsg(null, chunk));
            }
            return Flux.empty();
        });
        Flux<String> coordinatorResponse = Flux.defer(() -> {
            logger.info("Starting content processing");
            String fullContent = contentBuffer.toString();
            memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(fullContent, Map.of("agentId", Planner.AGENT_ID))));
            fullContent = message + "\n\n" + fullContent;
            return toCoordinator(fullContent, conversationId);

            /* Summary agent added between Planner and Coordinator was commented out due to increased processing time
               without significant improvement in results
            return summarizeContent(fullContent, conversationId)
                .flatMapMany(summary -> {
                    logger.info("Content summarization completed: {}", summary);
                    summary = message + "\n\n" + summary;
                    return toCoordinator(summary, conversationId);
                });
             */

        });
        
        return originalResponse
                .concatWith(coordinatorResponse)
                .filter(StringUtils::hasText);
    }
    
    /**
     * Summarizes content using the SummarizeAgent
     * @param content Original content to be summarized
     * @param conversationId Conversation identifier
     * @return Summarized content
     */
    private Mono<String> summarizeContent(String content, String conversationId) {
        logger.info("Beginning content summarization, original content length: {}", content.length());
        // Invoke SummarizeAgent to perform the summarization
        return Mono.fromCallable(() -> 
            agentManager.getAgent(SUMMARIZE_AGENT_ID).call(conversationId, content)
        );
    }
    
    Flux<String> toCoordinator(String content, String conversationId) {
        logger.info("User trip planning summary: {}", content);
        return agentManager.getAgent(COORDINATOR_AGENT_ID).stream(conversationId, content);
    }
}